# Copyright 2015 The Chromium Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""Decorator that validates XSRF tokens."""
from __future__ import print_function
from __future__ import division
from __future__ import absolute_import

import os
import six

from flask import abort, request

from google.appengine.ext import ndb

from dashboard.common import oauth2_utils
from dashboard.common import utils


class XsrfSecretKey(ndb.Model):
  """Stores a secret XSRF key for the site."""
  token = ndb.StringProperty(indexed=False)


def _ValidateToken(token, email):
  """Validates an XSRF token generated by GenerateXsrfToken."""
  return oauth2_utils.ValidateXsrfToken(
      _GetSecretKey(), token, user_id=email, action_id='')


def GenerateToken(email):
  return oauth2_utils.GenerateXsrfToken(
      _GetSecretKey(), user_id=email, action_id='')


def _GenerateNewSecretKey():
  """Returns a random XSRF secret key."""
  if six.PY2:
    return os.urandom(16).encode('hex')
  return os.urandom(16).hex()


def _GetSecretKey():
  """Gets or creates the secret key to use for validating XSRF tokens."""
  key_entity = ndb.Key('XsrfSecretKey', 'site').get()
  if not key_entity:
    key_entity = XsrfSecretKey(id='site', token=_GenerateNewSecretKey())
    key_entity.put()
  return key_entity.token.encode('ascii', 'ignore')


def TokenRequired(handler_method):
  """A decorator to require that the XSRF token be validated for the handler."""

  def CheckToken(self, *args, **kwargs):
    email = utils.GetEmail()
    token = str(self.request.get('xsrf_token'))
    if not email or not _ValidateToken(token, email):
      self.abort(403)
    handler_method(self, *args, **kwargs)

  return CheckToken


def TokenRequiredFlask(handler_method):
  """A decorator to require that the XSRF token be validated for the handler."""

  def CheckToken(*args, **kwargs):
    email = utils.GetEmail()
    token = str(request.values.get('xsrf_token'))
    if not email or not _ValidateToken(six.ensure_binary(token), email):
      abort(403)
    return handler_method(*args, **kwargs)

  return CheckToken
